import open3d as o3d
import numpy as np
import cv2
import pybullet as p
import pybullet_data

# --------------------------
# 1. Pybullet生成虚拟场景
# --------------------------
def generate_cube_soccer_scene():
    """生成立方体+足球的RGB-D数据"""
    print("Step 1/6：生成虚拟场景...")
    # 初始化仿真
    p.connect(p.GUI)
    p.setAdditionalSearchPath(pybullet_data.getDataPath())
    p.setGravity(0, 0, -9.8)
    
    # 加载场景(这里所有urdf文件是Pybullet下载时自带的)
    table = p.loadURDF("table/table.urdf", basePosition=[0, 0, 0])
    cube = p.loadURDF("cube.urdf", basePosition=[-0.2, 0, 0.775], globalScaling=0.1)
    soccerball = p.loadURDF("soccerball.urdf", basePosition=[0.2, 0, 0.767], globalScaling=0.2)
    p.changeVisualShape(cube, -1, rgbaColor=[0, 0, 1, 1])    # 蓝色立方体,"-1"指的是修改所有的Link
    p.changeVisualShape(soccerball, -1, rgbaColor=[0, 1, 0, 1])# 绿色足球
    
    # 相机参数
    camera_params = {
        "width": 640, "height": 480, "fov": 60, "near": 0.1, "far": 5.0,
        "eyePosition": [0, -1.5, 1.2], "targetPosition": [0, 0, 0.8], "upVector": [0, 0, 1]
    }
    proj_mat = p.computeProjectionMatrixFOV(
        camera_params["fov"], camera_params["width"]/camera_params["height"],
        camera_params["near"], camera_params["far"]
    )
    view_mat = p.computeViewMatrix(
        camera_params["eyePosition"], camera_params["targetPosition"], camera_params["upVector"]
    )
    
    # 渲染并保存RGB-D(拍照)
    rgb_array = p.getCameraImage(640, 480, view_mat, proj_mat, renderer=p.ER_BULLET_HARDWARE_OPENGL)[2]
    depth_array = p.getCameraImage(640, 480, view_mat, proj_mat, renderer=p.ER_BULLET_HARDWARE_OPENGL)[3]
    
    rgb_cv = cv2.cvtColor(rgb_array, cv2.COLOR_RGBA2BGR)
    depth_real = camera_params["near"] * camera_params["far"] / (camera_params["far"] - depth_array * (camera_params["far"] - camera_params["near"]))
    depth_uint16 = (depth_real * 1000).astype(np.uint16)#转成毫米
    
    cv2.imwrite("Soccer_Cube/data/cube_soccer_rgb.jpg", rgb_cv)
    cv2.imwrite("Soccer_Cube/data/cube_soccer_depth.png", depth_uint16)
    p.disconnect()
    print("场景数据保存完成：cube_soccer_rgb.jpg / cube_soccer_depth.png")
    return camera_params, rgb_cv, depth_uint16

# --------------------------
# 2. 从RGB-D生成点云
# --------------------------
def create_point_cloud(camera_params, rgb_cv, depth_uint16):
    """将RGB-D转换为3D点云"""
    print("Step 2/6：生成原始点云...")
    # 相机内参计算
    width, height = camera_params["width"], camera_params["height"]
    fx = fy = (width / 2) / np.tan(np.radians(camera_params["fov"] / 2))  # 焦距
    cx, cy = width / 2, height / 2  # 图像中心坐标
    intrinsic = o3d.camera.PinholeCameraIntrinsic(width, height, fx, fy, cx, cy)#针孔相机
    
    # 转换为Open3D RGBD格式
    o3d_rgb = o3d.geometry.Image(cv2.cvtColor(rgb_cv, cv2.COLOR_BGR2RGB))#BGR->RGB的Numpy数组
    o3d_depth = o3d.geometry.Image(depth_uint16)
    rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(#至关重要！2维->3维点云的桥梁
        o3d_rgb, o3d_depth, depth_scale=1000.0, depth_trunc=2.0, convert_rgb_to_intensity=False
    )
    
    # 生成点云并调整方向（防倒置）
    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intrinsic)#至关重要！从RGBD生成点云
    pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])  # 变换矩阵翻转YZ轴
    
    # 保存并可视化
    o3d.io.write_point_cloud("Soccer_Cube/data/raw_cube_soccer.ply", pcd)
    o3d.visualization.draw_geometries_with_editing([pcd], window_name="Step2：原始点云", width=800, height=600)
    print(f"原始点云保存完成（{len(pcd.points)}个点）：Soccer_Cube/data/raw_cube_soccer.ply")
    return pcd

# --------------------------
# 3. 点云预处理
# --------------------------
def preprocess_pcd(pcd):
    """优化点云质量：去除噪声+减少点数"""
    print("Step 3/6：点云预处理...")
    #  统计滤波（移除孤立点）
    _, inliers = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=1.0)#去噪后点云、噪声index
    denoised_pcd = pcd.select_by_index(inliers)#筛除噪声，保留原有点云结构
    
    #  下采样
    voxel_size = 0.005  # 5毫米体素（平衡精度和速度）
    downsampled_pcd = denoised_pcd.voxel_down_sample(voxel_size=voxel_size)
    
    # 保存并可视化
    o3d.io.write_point_cloud("Soccer_Cube/data/preprocessed_cube_soccer.ply", downsampled_pcd)
    #o3d.visualization.draw_geometries([downsampled_pcd], window_name="Step3：预处理后点云")
    o3d.visualization.draw_geometries_with_editing([downsampled_pcd], window_name="Step3：预处理后点云", width=800, height=600)
    print(f"预处理完成（{len(downsampled_pcd.points)}个点）：Soccer_Cube/data/preprocessed_cube_soccer.ply")
    return downsampled_pcd

# --------------------------
# 4. 几何分割
# --------------------------
def geometry_segmentation(pcd):
    """用平面分割+DBSCAN聚类：桌面（0）、立方体（1）、足球（2）"""
    print("Step 4/6：几何分割（平面+聚类）...")
    
    #-----平面分割：通过法向量筛选桌面-----
    max_attempts = 10 
    desktop_inliers = []
    found_valid_plane = False
    for _ in range(max_attempts):
        plane_model, inliers = pcd.segment_plane(#切割最有可能的平面
            distance_threshold=0.015,  # 点到平面最大距离，即可不平整
            ransac_n=3,#ax+by+cz+d=0方程参数数量
            num_iterations=1000
        )
        # 计算平面法向量（plane_model格式：[a, b, c, d]，对应ax+by+cz+d=0）
        normal = np.asarray(plane_model)[:3]
        normal = normal / np.linalg.norm(normal)  # 归一化法向量
        # 点积判定与[0,0,1]相似度
        dot_product = np.dot(normal, [0, 0, 1])
        if dot_product > 0.9:  # 相似度足够高（接近垂直向上）
            found_valid_plane = True
            desktop_inliers = inliers
            break
    # 避免中断，取第一次分割结果保底
    if not found_valid_plane:
        plane_model, desktop_inliers = pcd.segment_plane(
            distance_threshold=0.01,
            ransac_n=3,
            num_iterations=1000
        )
    table_pcd = pcd.select_by_index(desktop_inliers)  # 桌面点云
    objects_pcd = pcd.select_by_index(desktop_inliers, invert=True)  # 物体点云（立方体+足球）
    
    #-----DBSCAN聚类：区分立方体和足球（基于空间距离）------
    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug):
        labels = np.array(objects_pcd.cluster_dbscan(eps=0.025, min_points=30, print_progress=True))
    
    #-----标签映射：确保标签≤2-----
    max_label = labels.max()
    print(f"DBSCAN聚类结果：共{max_label + 1}个物体簇（含噪声）")
    
    object_labels = labels + 1  # 初始映射：0→1，1→2，-1（噪声）→0
    object_labels[object_labels > 2] = 0  # 超过2的标签视为噪声（归为0）
    
    #-----合并桌面和物体的标签-----
    table_labels = np.zeros(len(table_pcd.points), dtype=np.int32)
    combined_pcd = table_pcd + objects_pcd
    combined_labels = np.concatenate([table_labels, object_labels])#整合分离的点云，便于操作
    
    #-----可视化分割结果-----
    color_map = np.array([[0.5, 0.5, 0.5], [0, 1, 0], [0, 0, 1]]) 
    safe_labels = np.clip(combined_labels, 0, 2)  # 强制标签在0-2
    seg_colors = color_map[safe_labels]  # 分配颜色

    #-----手动复制点云数据-----（没搞懂为啥复制函数用不了，但手动也行）
    seg_pcd = o3d.geometry.PointCloud()  # 创建新的点云对象
    seg_pcd.points = combined_pcd.points  # 复制点坐标
    seg_pcd.colors = o3d.utility.Vector3dVector(seg_colors)  # 赋值新颜色

    # 显示窗口（兼容模式）
    o3d.visualization.draw_geometries_with_editing(
        [seg_pcd], 
        window_name="Step4：语义分割结果（灰=桌面，蓝=立方体，绿=足球）",
        width=800, height=600
    )
    
    # 打印标签分布（验证是否只有0、1、2）
    label_counts = np.bincount(combined_labels)
    print("语义分割完成，标签分布：")
    print(f"- 桌面（标签0）：{label_counts[0]}个点")
    print(f"- 立方体（标签1）：{label_counts[1]}个点" if len(label_counts)>1 else "- 立方体（标签1）：0个点")
    print(f"- 足球（标签2）：{label_counts[2]}个点" if len(label_counts)>2 else "- 足球（标签2）：0个点")
    
    return combined_pcd, combined_labels

# --------------------------
# 5. 提取立方体+足球点云
# --------------------------
def extract_targets(pcd, mapped_labels):
    """从分割结果中提取立方体和足球点云"""
    print("Step 5/6：提取目标点云...")
    # 按标签筛选
    cube_indices = np.where(mapped_labels == 1)[0]  # 立方体标签=1
    soccer_indices = np.where(mapped_labels == 2)[0]  # 足球标签=2
    
    cube_pcd = pcd.select_by_index(cube_indices)
    soccer_pcd = pcd.select_by_index(soccer_indices)
    
    # 上色（确保可视化清晰，与分割颜色一致）
    cube_pcd.paint_uniform_color([0, 1, 0])    
    soccer_pcd.paint_uniform_color([0, 0, 1])  
    
    # 可视化提取结果
    o3d.visualization.draw_geometries(
        [cube_pcd, soccer_pcd], 
        window_name="Step5：提取的目标（蓝=立方体，绿=足球）"
    )
    print(f"目标提取完成：立方体({len(cube_pcd.points)}个点)，足球({len(soccer_pcd.points)}个点)")
    return cube_pcd, soccer_pcd

# --------------------------
# 6. 保存带语义标签的PLY文件
# --------------------------
def save_ply(cube_pcd, soccer_pcd):
    """保存最终结果，用颜色通道存储语义标签(当前Open3d版本不支持attributes储存自定义标签)"""
    print("Step 6/6：保存PLY文件...")
    
    # 为立方体添加语义标签（1）：用红色通道存储（值=1.0代表标签1）
    # 创建形状为(N, 3)的颜色数组，红色通道=1，其他通道=0
    cube_label_colors = np.zeros((len(cube_pcd.points), 3), dtype=np.float64)
    cube_label_colors[:, 0] = 1.0  # 红色通道=1 → 标签1（立方体）
    cube_pcd.colors = o3d.utility.Vector3dVector(cube_label_colors)
    
    # 为足球添加语义标签（2）：用绿色通道存储（值=1.0代表标签2）
    soccer_label_colors = np.zeros((len(soccer_pcd.points), 3), dtype=np.float64)
    soccer_label_colors[:, 1] = 1.0  # 绿色通道=1 → 标签2（足球）
    soccer_pcd.colors = o3d.utility.Vector3dVector(soccer_label_colors)
    
    # 保存点云（颜色通道会被写入PLY文件）
    o3d.io.write_point_cloud("Soccer_Cube/data/cube_final.ply", cube_pcd)
    o3d.io.write_point_cloud("Soccer_Cube/data/soccer_final.ply", soccer_pcd)
    print("最终文件保存完成（语义标签存储在颜色通道）：")
    print("- 立方体（标签1）：蓝色通道=1.0 → cube_final.ply")
    print("- 足球（标签2）：绿色通道=1.0 → soccer_final.ply")


# --------------------------
# 主函数
# --------------------------
if __name__ == "__main__":
    # 1. 生成场景
    camera_params, rgb_cv, depth_uint16 = generate_cube_soccer_scene()
    # 2. 生成点云
    raw_pcd = create_point_cloud(camera_params, rgb_cv, depth_uint16)
    # 3. 预处理
    preprocessed_pcd = preprocess_pcd(raw_pcd)
    # 4. 几何分割
    seg_pcd, labels = geometry_segmentation(preprocessed_pcd)
    # 5. 提取目标
    cube_pcd, soccer_pcd = extract_targets(seg_pcd, labels)
    # 6. 保存结果
    save_ply(cube_pcd, soccer_pcd)
    print("\n全流程完成！可在当前目录查看生成的PLY文件。")